{
 "cells": [
  {
   "cell_type": "raw",
   "id": "4b1f2015915881ca",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "##读入数据集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4d3f1b4741913ce",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-06-20T08:06:56.361921Z",
     "start_time": "2024-06-20T08:06:56.100644800Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "标签到ID的映射字典：\n",
      "{1: 0, 2: 1, 3: 2, 4: 3, 5: 4}\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "import pandas as pd\n",
    "from sklearn.model_selection import train_test_split\n",
    "# 读取Excel文件\n",
    "df = pd.read_excel('data.xlsx')\n",
    "# 提取所有标签\n",
    "tags = df['评价星级']\n",
    "# 获取唯一标签集合\n",
    "unique_tags = set(tags)\n",
    "\n",
    "# 创建标签到ID的映射字典\n",
    "label2id_dict = {tag: idx for idx, tag in enumerate(unique_tags)}\n",
    "\n",
    "print(\"标签到ID的映射字典：\")\n",
    "print(label2id_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "2c7869cddebb6877",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-06-20T08:06:56.422219900Z",
     "start_time": "2024-06-20T08:06:56.366612700Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5673\n",
      "3782\n"
     ]
    }
   ],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "X = df['评价内容']\n",
    "y = df['评价星级']\n",
    "# 将数据划分为训练集和验证集\n",
    "X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.4, random_state=12)\n",
    "\n",
    "X_train = X_train.tolist() if isinstance(X_train, pd.Series) else X_train\n",
    "y_train = y_train.tolist() if isinstance(y_train, pd.Series) else y_train\n",
    "\n",
    "X_val = X_val.tolist() if isinstance(X_val, pd.Series) else X_val\n",
    "y_val = y_val.tolist() if isinstance(y_val, pd.Series) else y_val\n",
    "\n",
    "print(len(X_train))\n",
    "print(len(X_val))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "21abf144-8748-4afc-af04-1ed6c200c93e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-06-20T08:07:02.448627500Z",
     "start_time": "2024-06-20T08:06:56.397827600Z"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from transformers import BertTokenizer, ErnieForSequenceClassification\n",
    "\n",
    "MAX_LENGTH=128\n",
    "\n",
    "class CustomDataset(Dataset):\n",
    "    def __init__(self, texts,tags,tokenizer,label2id, max_length):\n",
    "        self.texts = texts\n",
    "        self.tags = tags\n",
    "        self.max_length = max_length\n",
    "        self.tokenizer = tokenizer\n",
    "        self.label2id = label2id\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.texts)\n",
    "    \n",
    "    def find_split_point(self,text,start,max_length):\n",
    "        split_point = min(start+max_length,len(text)-1)\n",
    "        while split_point>start:\n",
    "            if text[split_point] in \"。！？：；\": \n",
    "                if split_point-start+1<=max_length:\n",
    "                    return split_point+1\n",
    "            split_point-=1\n",
    "        return start+max_length\n",
    "    \n",
    "    def segment_and_tokenize(self,text,tag):\n",
    "        input_ids, attention_masks, labels = [], [], []\n",
    "        i = 0\n",
    "        while i<len(text):\n",
    "            if len(text[i:])>self.max_length-2:\n",
    "                split_point = self.find_split_point(text,i,self.max_length-2)\n",
    "            else:\n",
    "                split_point=len(text)\n",
    "            segment_text = text[i:split_point]\n",
    "            encoded_dict = self.tokenizer.encode_plus(\n",
    "                segment_text,\n",
    "                is_split_into_words=True,\n",
    "                max_length=self.max_length,\n",
    "                padding='max_length',\n",
    "                truncation=True,\n",
    "                return_attention_mask=True\n",
    "            )\n",
    "            # 添加到输出列表\n",
    "            input_ids.append(encoded_dict['input_ids'])\n",
    "            attention_masks.append(encoded_dict['attention_mask'])\n",
    "            labels.append(self.label2id.get(tag))\n",
    "            i = split_point  # 更新索引，继绀处理下一个段落\n",
    "        input_ids = torch.tensor(input_ids, dtype=torch.long)\n",
    "        attention_masks = torch.tensor(attention_masks, dtype=torch.long)\n",
    "        labels = torch.tensor(labels, dtype=torch.long)\n",
    "    \n",
    "        return {\n",
    "            'input_ids': input_ids,\n",
    "            'attention_mask': attention_masks,\n",
    "            'labels': labels\n",
    "        }\n",
    "            \n",
    "    def __getitem__(self, idx):\n",
    "        \"\"\" 获取指定索引处的数据项 \"\"\"\n",
    "        text, tag = self.texts[idx], self.tags[idx]\n",
    "        # 进行分段和编码，每个段都填充到最大长度\n",
    "        return self.segment_and_tokenize(text, tag)\n",
    "    \n",
    "    def collate_fn(self, batch):\n",
    "        \"\"\"\n",
    "        数据集的自定义批处理函数。\n",
    "        将多个元素合并成一个批次，注意它需要对好几个段拼接在一起的情况进行处理。\n",
    "\n",
    "        Args:\n",
    "            batch (list): 包含多个数据项的列表。\n",
    "\n",
    "        Returns:\n",
    "            dict: 包含批量input_ids, attention_mask, 和labels的字典。\n",
    "        \"\"\"\n",
    "        # 给定的batch是一个字典列表（每个元素都是__getitem__的输出）\n",
    "        input_ids = torch.cat([item['input_ids'] for item in batch], dim=0)\n",
    "        attention_mask = torch.cat([item['attention_mask'] for item in batch], dim=0)\n",
    "        labels = torch.cat([item['labels'] for item in batch], dim=0)\n",
    "\n",
    "        return {\n",
    "            'input_ids': input_ids,\n",
    "            'attention_mask': attention_mask,\n",
    "            'labels': labels\n",
    "        }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "4551abf2-8997-4b18-a5ec-f284df528035",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-06-20T08:07:43.720514100Z",
     "start_time": "2024-06-20T08:07:43.488475200Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/miniconda3/lib/python3.8/site-packages/torch/_utils.py:776: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
      "  return self.fget.__get__(instance, owner)()\n",
      "Some weights of ErnieForSequenceClassification were not initialized from the model checkpoint at ernie-3.0-mini-zh and are newly initialized: ['classifier.bias', 'classifier.weight', 'ernie.pooler.dense.bias', 'ernie.pooler.dense.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5673\n",
      "Training data:\n",
      "{'input_ids': tensor([[   1, 1470, 1346,  ...,    0,    0,    0],\n",
      "        [   1, 1082,  185,  ...,    0,    0,    0],\n",
      "        [   1,  170,   65,  ...,    0,    0,    0],\n",
      "        ...,\n",
      "        [   1,  247,   98,  ...,    0,    0,    0],\n",
      "        [   1,  287,   40,  ...,    0,    0,    0],\n",
      "        [   1,  416,   50,  ...,    0,    0,    0]]), 'attention_mask': tensor([[1, 1, 1,  ..., 0, 0, 0],\n",
      "        [1, 1, 1,  ..., 0, 0, 0],\n",
      "        [1, 1, 1,  ..., 0, 0, 0],\n",
      "        ...,\n",
      "        [1, 1, 1,  ..., 0, 0, 0],\n",
      "        [1, 1, 1,  ..., 0, 0, 0],\n",
      "        [1, 1, 1,  ..., 0, 0, 0]]), 'labels': tensor([4, 4, 2, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 1, 4, 4,\n",
      "        4, 4, 4, 4, 4, 4, 4, 4])}\n",
      "Decoded Sentence: 屏 幕 效 果 ： 优 异 ， 散 热 性 能 ： 优 异 ， 外 形 外 观 ： 质 感 ， 轻 薄 程 度 ： 完 美 ， 单 手 开 合 ， 重 量 适 中 ， 硬 件 配 置 也 不 拖 后 腿 ， 轻 薄 本 中 非 常 能 打 的 一 款 ， 相 信 在 后 续 使 用 体 验 中 也 非 常 棒 ！\n",
      "Input IDs: tensor([    1,  1470,  1346,   344,   228,    74,   336,   712,     4,   851,\n",
      "          506,    92,    52,    74,   336,   712,     4,   137,   193,   137,\n",
      "          394,    74,   207,   345,     4,   709,  1378,   163,   130,    74,\n",
      "          328,   188,     4,   269,   247,    88,    95,     4,   118,   150,\n",
      "          582,    12,     4,  1003,   259,   483,   521,   105,    16,  1950,\n",
      "           49,  1886,     4,   709,  1378,    89,    12,   465,   223,    52,\n",
      "          445,     5,     7,   699,     4,   156,   212,    11,    49,   607,\n",
      "          175,    29,    82,   412,    12,   105,   465,   223,  2203, 12044,\n",
      "            2,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0])\n",
      "Attention Mask: tensor([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
      "        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
      "        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
      "        1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "        0, 0, 0, 0, 0, 0, 0, 0])\n",
      "Labels: tensor(4)\n",
      "\n",
      "Validation data:\n",
      "Decoded Sentence: 非 常 好. 建 议 大 家 购 买. 客 服 很 耐 心. 美 中 不 足 的 是 物 流 太 慢 了. ， 屏 幕 效 果 ： 高 清 ， 运 行 速 度 ： 快 ， 轻 薄 程 度 ： 非 常 轻\n",
      "Input IDs: tensor([   1,  465,  223,  170,   42,   81,  454,   19,   50,  817, 1042,   42,\n",
      "         478,  231,  321, 1275,  111,   42,  188,   12,   16,  581,    5,   10,\n",
      "         122,  235,  512, 1136,   15,   42,    4, 1470, 1346,  344,  228,   74,\n",
      "          69,  334,    4,  287,   40,  537,  130,   74,  532,    4,  709, 1378,\n",
      "         163,  130,   74,  465,  223,  709,    2,    0,    0,    0,    0,    0,\n",
      "           0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,\n",
      "           0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,\n",
      "           0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,\n",
      "           0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,\n",
      "           0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,\n",
      "           0,    0,    0,    0,    0,    0,    0,    0])\n",
      "Attention Mask: tensor([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
      "        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
      "        1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "        0, 0, 0, 0, 0, 0, 0, 0])\n",
      "Labels: tensor(4)\n"
     ]
    }
   ],
   "source": [
    "# 加载预训练的ERNIE模型和tokenizer\n",
    "model_name = \"ernie-3.0-mini-zh\" \n",
    "tokenizer = BertTokenizer.from_pretrained(model_name)\n",
    "model = ErnieForSequenceClassification.from_pretrained(model_name,num_labels=len(label2id_dict))\n",
    "# 创建训练集和验证集的 Dataset\n",
    "train_dataset = CustomDataset(X_train, y_train,tokenizer,label2id_dict,MAX_LENGTH)\n",
    "val_dataset = CustomDataset(X_val, y_val,tokenizer,label2id_dict,MAX_LENGTH)\n",
    "print(len(train_dataset))\n",
    "train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True,collate_fn=train_dataset.collate_fn)\n",
    "dev_loader = DataLoader(val_dataset, batch_size=32,collate_fn=val_dataset.collate_fn)\n",
    "# 迭代训练数据加载器，打印第一批数据\n",
    "print(\"Training data:\")\n",
    "for batch in train_loader:\n",
    "    print(batch)\n",
    "    input_ids = batch['input_ids']\n",
    "    attention_mask = batch['attention_mask']\n",
    "    labels = batch['labels']\n",
    "    # 使用tokenizer.decode将第一个input_ids转换回文本字符串\n",
    "    decoded_sentence = tokenizer.decode(input_ids[0], skip_special_tokens=True)\n",
    "    \n",
    "    print(\"Decoded Sentence:\", decoded_sentence)\n",
    "    print(\"Input IDs:\", input_ids[0])\n",
    "    print(\"Attention Mask:\", attention_mask[0])\n",
    "    print(\"Labels:\", labels[0])\n",
    "    break\n",
    "\n",
    "# 迭代验证数据加载器，打印第一批数据\n",
    "print(\"\\nValidation data:\")\n",
    "for batch in dev_loader:\n",
    "    input_ids = batch['input_ids']\n",
    "    attention_mask = batch['attention_mask']\n",
    "    labels = batch['labels']\n",
    "    # 使用tokenizer.decode将第一个input_ids转换回文本字符串\n",
    "    decoded_sentence = tokenizer.decode(input_ids[0], skip_special_tokens=True)\n",
    "    print(\"Decoded Sentence:\", decoded_sentence)\n",
    "    print(\"Input IDs:\", input_ids[0])\n",
    "    print(\"Attention Mask:\", attention_mask[0])\n",
    "    print(\"Labels:\", labels[0])\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "761cacd4-493e-4e78-9473-316869241726",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-06-20T08:19:09.025852900Z",
     "start_time": "2024-06-20T08:16:20.340728300Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/miniconda3/lib/python3.8/site-packages/transformers/optimization.py:588: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1 done - Train Loss: 0.4511 - Train Acc: 0.8751\n",
      "Val Loss: 0.3349 - Val Acc: 0.8997\n",
      "Epoch 2 done - Train Loss: 0.3260 - Train Acc: 0.9034\n",
      "Val Loss: 0.3493 - Val Acc: 0.9037\n",
      "Epoch 3 done - Train Loss: 0.2847 - Train Acc: 0.9127\n",
      "Val Loss: 0.3256 - Val Acc: 0.9083\n",
      "Epoch 4 done - Train Loss: 0.2438 - Train Acc: 0.9233\n",
      "Val Loss: 0.3355 - Val Acc: 0.9050\n",
      "Epoch 5 done - Train Loss: 0.2081 - Train Acc: 0.9330\n",
      "Val Loss: 0.3526 - Val Acc: 0.9072\n",
      "Epoch 6 done - Train Loss: 0.1733 - Train Acc: 0.9453\n",
      "Val Loss: 0.3695 - Val Acc: 0.9032\n",
      "Epoch 7 done - Train Loss: 0.1449 - Train Acc: 0.9521\n",
      "Val Loss: 0.3849 - Val Acc: 0.8990\n",
      "Epoch 8 done - Train Loss: 0.1204 - Train Acc: 0.9614\n",
      "Val Loss: 0.4155 - Val Acc: 0.9022\n",
      "Epoch 9 done - Train Loss: 0.1007 - Train Acc: 0.9692\n",
      "Val Loss: 0.4289 - Val Acc: 0.8937\n",
      "Epoch 10 done - Train Loss: 0.0907 - Train Acc: 0.9737\n",
      "Val Loss: 0.4295 - Val Acc: 0.8972\n"
     ]
    }
   ],
   "source": [
    "from transformers import AdamW, get_scheduler\n",
    "\n",
    "optimizer = AdamW(model.parameters(), lr=5e-5)\n",
    "num_epochs = 10\n",
    "num_training_steps = num_epochs * len(train_loader)\n",
    "lr_scheduler = get_scheduler(\n",
    "    \"linear\", \n",
    "    optimizer=optimizer, \n",
    "    num_warmup_steps=0, \n",
    "    num_training_steps=num_training_steps\n",
    ")\n",
    "\n",
    "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "model.to(device)\n",
    "best_val_accuracy = 0.0\n",
    "for epoch in range(num_epochs):\n",
    "    model.train()\n",
    "    total_train_loss = total_train_accuracy = 0\n",
    "    for batch in train_loader:\n",
    "        batch = {k: v.to(device) for k, v in batch.items()}\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(**batch)\n",
    "        loss = outputs.loss\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        lr_scheduler.step()\n",
    "\n",
    "        total_train_loss += loss.item()\n",
    "        logits = outputs.logits\n",
    "        predictions = torch.argmax(logits, dim=-1)\n",
    "        evaluator = (predictions == batch['labels']) \n",
    "        total_train_accuracy += evaluator.sum().item() / evaluator.numel()\n",
    "    # 计算平均训练损失和准确率\n",
    "    avg_train_loss = total_train_loss / len(train_loader)\n",
    "    avg_train_accuracy = total_train_accuracy / len(train_loader)\n",
    "    log_str = f\"Epoch {epoch+1} done - Train Loss: {avg_train_loss:.4f} - Train Acc: {avg_train_accuracy:.4f}\"\n",
    "    print(log_str)\n",
    "    # 验证循环\n",
    "    model.eval()\n",
    "    total_val_loss = total_val_accuracy = 0\n",
    "    total_loss = 0\n",
    "    with torch.no_grad():\n",
    "        for batch in dev_loader:\n",
    "            batch = {k: v.to(device) for k, v in batch.items()}\n",
    "            outputs = model(**batch)\n",
    "            loss = outputs.loss\n",
    "            total_val_loss += loss.item()\n",
    "\n",
    "            logits = outputs.logits\n",
    "            predictions = torch.argmax(logits, dim=-1)\n",
    "            evaluator = (predictions == batch['labels']) \n",
    "            total_val_accuracy += evaluator.sum().item() / evaluator.numel()\n",
    "    avg_val_loss = total_val_loss / len(dev_loader)\n",
    "    avg_val_accuracy = total_val_accuracy / len(dev_loader)\n",
    "    log_str = f\"Val Loss: {avg_val_loss:.4f} - Val Acc: {avg_val_accuracy:.4f}\"\n",
    "    print(log_str)\n",
    "    torch.save(model.state_dict(),'model.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "81739e25-d00b-485b-9065-d93776c3e401",
   "metadata": {
    "ExecuteTime": {
     "start_time": "2024-06-20T08:07:03.031910Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "def predict(texts, model, tokenizer, label2id, max_length=128, batch_size=32, device='cuda'):\n",
    "    # 创建临时数据集\n",
    "    class TempDataset(Dataset):\n",
    "        def __init__(self, texts, tokenizer, label2id, max_length):\n",
    "            self.texts = texts\n",
    "            self.tokenizer = tokenizer\n",
    "            self.max_length = max_length\n",
    "            self.label2id = label2id\n",
    "        \n",
    "        def __len__(self):\n",
    "            return len(self.texts)\n",
    "        \n",
    "        def find_split_point(self, text, start, max_length):\n",
    "            split_point = min(start + max_length, len(text) - 1)\n",
    "            while split_point > start:\n",
    "                if text[split_point] in \"。！？：；\":\n",
    "                    if split_point - start + 1 <= max_length:\n",
    "                        return split_point + 1\n",
    "                split_point -= 1\n",
    "            return start + max_length\n",
    "        \n",
    "        def segment_and_tokenize(self, text):\n",
    "            input_ids, attention_masks = [], []\n",
    "            i = 0\n",
    "            while i < len(text):\n",
    "                if len(text[i:]) > self.max_length - 2:\n",
    "                    split_point = self.find_split_point(text, i, self.max_length - 2)\n",
    "                else:\n",
    "                    split_point = len(text)\n",
    "                segment_text = text[i:split_point]\n",
    "                encoded_dict = self.tokenizer.encode_plus(\n",
    "                    segment_text,\n",
    "                    is_split_into_words=True,\n",
    "                    max_length=self.max_length,\n",
    "                    padding='max_length',\n",
    "                    truncation=True,\n",
    "                    return_attention_mask=True\n",
    "                )\n",
    "                input_ids.append(encoded_dict['input_ids'])\n",
    "                attention_masks.append(encoded_dict['attention_mask'])\n",
    "                i = split_point\n",
    "        \n",
    "            input_ids = torch.tensor(input_ids, dtype=torch.long)\n",
    "            attention_masks = torch.tensor(attention_masks, dtype=torch.long)\n",
    "    \n",
    "            return {\n",
    "                'input_ids': input_ids,\n",
    "                'attention_mask': attention_masks\n",
    "            }\n",
    "        \n",
    "        def __getitem__(self, idx):\n",
    "            text = self.texts[idx]\n",
    "            return self.segment_and_tokenize(text)\n",
    "        \n",
    "        def collate_fn(self, batch):\n",
    "            input_ids = torch.cat([item['input_ids'] for item in batch], dim=0)\n",
    "            attention_mask = torch.cat([item['attention_mask'] for item in batch], dim=0)\n",
    "    \n",
    "            return {\n",
    "                'input_ids': input_ids,\n",
    "                'attention_mask': attention_mask\n",
    "            }\n",
    "    \n",
    "    # 创建临时数据集和 DataLoader\n",
    "    temp_dataset = TempDataset(texts, tokenizer, label2id, max_length)\n",
    "    temp_loader = DataLoader(temp_dataset, batch_size=batch_size, collate_fn=temp_dataset.collate_fn)\n",
    "    \n",
    "    model.to(device)\n",
    "    model.eval()\n",
    "    \n",
    "    predictions = []\n",
    "    with torch.no_grad():\n",
    "        for batch in temp_loader:\n",
    "            input_ids = batch['input_ids'].to(device)\n",
    "            attention_mask = batch['attention_mask'].to(device)\n",
    "            outputs = model(input_ids, attention_mask=attention_mask)\n",
    "            logits = outputs.logits\n",
    "            preds = torch.argmax(logits, dim=-1)\n",
    "            predictions.extend(preds.cpu().numpy())\n",
    "    \n",
    "    id2label = {v: k for k, v in label2id.items()}\n",
    "    predicted_labels = [id2label[pred] for pred in predictions]\n",
    "    \n",
    "    return predicted_labels\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "9a36a69a-bb4b-4495-b46b-d061d876b1f3",
   "metadata": {
    "ExecuteTime": {
     "start_time": "2024-06-20T08:07:03.033910700Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'texts' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[17], line 6\u001b[0m\n\u001b[1;32m      4\u001b[0m output_file \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mpredictions.txt\u001b[39m\u001b[38;5;124m'\u001b[39m\n\u001b[1;32m      5\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mopen\u001b[39m(output_file, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mw\u001b[39m\u001b[38;5;124m'\u001b[39m, encoding\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mutf-8\u001b[39m\u001b[38;5;124m'\u001b[39m) \u001b[38;5;28;01mas\u001b[39;00m file:\n\u001b[0;32m----> 6\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m text, label \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(\u001b[43mtexts\u001b[49m, predictions):\n\u001b[1;32m      7\u001b[0m         file\u001b[38;5;241m.\u001b[39mwrite(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mtext\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;130;01m\\t\u001b[39;00m\u001b[38;5;132;01m{\u001b[39;00mlabel\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m      9\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mPredictions saved to \u001b[39m\u001b[38;5;132;01m{\u001b[39;00moutput_file\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m'\u001b[39m)\n",
      "\u001b[0;31mNameError\u001b[0m: name 'texts' is not defined"
     ]
    }
   ],
   "source": [
    "texts = ['成分很安全的一款面膜，对香精敏感的美眉慎入，作为日常补水款足够了。', '滋润，不油腻，味道也喜欢。', '一直在用，感觉还是不错，这次还有小礼品，谢谢啦！', '还有快递小哥，今天下了雪如期送达，感谢！']\n",
    "predictions = predict(texts, model, tokenizer, label2id_dict, max_length=MAX_LENGTH, device='cuda')\n",
    "# 输出预测结果\n",
    "output_file = 'predictions.txt'\n",
    "with open(output_file, 'w', encoding='utf-8') as file:\n",
    "    for text, label in zip(texts, predictions):\n",
    "        print(f'{text}\\t{label}\\n\\n')\n",
    "\n",
    "print(f'Predictions saved to {output_file}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c6888b95e6adc94",
   "metadata": {
    "ExecuteTime": {
     "start_time": "2024-06-20T08:07:03.035947600Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
